// Date:   Thu Aug 14 12:14:39 AM 2025
// Mail:   lunar_ubuntu@qq.com
// Author: https://github.com/xiaoqixian

#include "common.hpp"
#include <cstddef>
#include <random>
#include <libavutil/frame.h>
#include <cmath>
#include <cstring>
#include <vector>
#include <algorithm>
#include <openvino/openvino.hpp>

extern "C" {
#include <libswscale/swscale.h>
#include <libavutil/pixfmt.h>
#include <libavutil/imgutils.h>
}

ov::Tensor mat_to_tensor(const cv::Mat& img, int target_w, int target_h) {
  // 1. resize
  cv::Mat resized;
  cv::resize(img, resized, cv::Size(target_w, target_h));

  // 2. BGR -> RGB
  cv::cvtColor(resized, resized, cv::COLOR_BGR2RGB);

  // 3. convert to float and normalize [0,1]
  resized.convertTo(resized, CV_32F, 1.0 / 255.0);

  // 4. create tensor
  ov::Shape shape = {3, (size_t)target_h, (size_t)target_w};
  ov::Tensor tensor(ov::element::f32, shape);
  float* tensor_data = tensor.data<float>();

  // 5. HWC -> CHW
  std::vector<cv::Mat> chw(3);
  for (int c = 0; c < 3; ++c) {
    chw[c] = cv::Mat(target_h, target_w, CV_32F, tensor_data + c * target_h * target_w);
  }
  cv::split(resized, chw);

  return tensor;
}

float frame_diff_ratio(cv::Mat const& img1, cv::Mat const& img2) {
  // Ensure same size and type
  if (img1.size() != img2.size() || img1.type() != img2.type()) {
    // throw std::invalid_argument("Images must have the same size and type.");
    return 1.0;
  }

  cv::Mat diff;
  cv::absdiff(img1, img2, diff);

  // Convert to grayscale to simplify difference
  if (diff.channels() > 1) {
    cv::cvtColor(diff, diff, cv::COLOR_BGR2GRAY);
  }

  // Threshold to get binary difference mask
  cv::Mat diff_mask;
  cv::threshold(diff, diff_mask, 30, 255, cv::THRESH_BINARY); // 30 可调

  // Calculate ratio
  float diff_pixels = static_cast<float>(cv::countNonZero(diff_mask));
  float total_pixels = static_cast<float>(diff_mask.total());

  return diff_pixels / total_pixels;
}

int weighted_random_sample(std::vector<std::size_t> const& weights) {
  static std::random_device rd;
  static std::mt19937 gen(rd());

  std::vector<std::size_t> prefix(weights.size());
  std::partial_sum(weights.begin(), weights.end(), prefix.begin());
  auto const total = prefix.back();

  std::uniform_int_distribution<std::size_t> dist(1, total);
  auto const rnd = dist(gen);

  auto it = std::lower_bound(prefix.begin(), prefix.end(), rnd);
  return static_cast<int>(it - prefix.begin());
}

void mat_preprocess(const cv::Mat& img_bgr, const std::size_t img_size, ov::Tensor& tensor) {
  // === BGR -> RGB ===
  cv::Mat img;
  cv::cvtColor(img_bgr, img, cv::COLOR_BGR2RGB);

  // === Resize ===
  const int target_size = static_cast<int>(img_size * 1.14); // 255
  int h = img.rows;
  int w = img.cols;
  int new_h, new_w;
  if (h < w) {
    new_h = target_size;
    new_w = static_cast<int>(w * target_size / h);
  } else {
    new_h = static_cast<int>(h * target_size / w);
    new_w = target_size;
  }
  cv::resize(img, img, cv::Size(new_w, new_h), 0, 0, cv::INTER_LINEAR);

  // === CenterCrop 224 ===
  int start_h = (img.rows - img_size) / 2;
  int start_w = (img.cols - img_size) / 2;
  cv::Rect crop_roi(start_w, start_h, img_size, img_size);
  img = img(crop_roi).clone();

  // === Convert to float32 and normalize ===
  img.convertTo(img, CV_32F, 1.0 / 255.0);

  // Mean and std
  const std::vector<float> mean = {0.485f, 0.456f, 0.406f};
  const std::vector<float> std  = {0.229f, 0.224f, 0.225f};

  // HWC -> CHW
  std::vector<cv::Mat> channels(3);
  cv::split(img, channels);
  for (size_t i = 0; i < 3; i++) {
    channels[i] = (channels[i] - mean[i]) / std[i];
  }

  // Copy data to ov::Tensor
  float* tensor_data = tensor.data<float>();
  size_t channel_size = img_size * img_size;
  for (size_t c = 0; c < 3; c++) {
    memcpy(tensor_data + c * channel_size,
         channels[c].data,
         channel_size * sizeof(float));
  }
}

void av_preprocess(AVFrame* frame, std::size_t input_size, ov::Tensor& tensor) {
  const int target_size = static_cast<int>(input_size * 1.14); // 255
  const int h = frame->height;
  const int w = frame->width;
  int new_h, new_w;
  if (h < w) {
    new_h = target_size;
    new_w = static_cast<int>(w * target_size / h);
  } else {
    new_h = static_cast<int>(h * target_size / w);
    new_w = target_size;
  }

  constexpr float mean[] = {0.5, 0, 0};
  constexpr float stdd[] = {0.5, 1, 1};
  float* ov_data = tensor.data<float>();

  for (int ch = 0; ch < 3; ch++) {
    int real_h = h, real_w = w;
    switch (frame->format) {
      case AV_PIX_FMT_YUV420P:
        real_h = h / 2;
        real_w = w / 2;
        break;
      case AV_PIX_FMT_YUV422P:
        real_w = w / 2;
        break;
      case AV_PIX_FMT_YUV444P:
        break;
      default:
        throw std::runtime_error(cv::format("unexpected frame format: %d\n", frame->format));
    }

    cv::Mat orig(real_h, real_w, CV_8UC1, frame->data[ch], frame->linesize[ch]);
    cv::Mat mat;
    cv::resize(orig, mat, cv::Size(new_w, new_h), 0, 0, cv::INTER_LINEAR);
    const int start_h = (new_h - input_size) / 2;
    const int start_w = (new_w - input_size) / 2;
    cv::Rect crop_roi(start_w, start_h, input_size, input_size);
    mat = mat(crop_roi).clone();
    mat.convertTo(mat, CV_32F, 1.0/255.0);
    mat = (mat - mean[ch]) / stdd[ch];

    // copy to ov::Tensor
    std::memcpy(ov_data, mat.data, input_size * input_size * sizeof(float));
    ov_data += input_size * input_size;

  }
}

// Simple bilinear resize for RGB image
static void bilinear_resize(const uint8_t* src, int src_w, int src_h,
                     uint8_t* dst, int dst_w, int dst_h) {
  for (int y = 0; y < dst_h; y++) {
    float fy = (y + 0.5f) * src_h / dst_h - 0.5f;
    int y0 = std::max(0, (int)std::floor(fy));
    int y1 = std::min(src_h - 1, y0 + 1);
    float wy = fy - y0;

    for (int x = 0; x < dst_w; x++) {
      float fx = (x + 0.5f) * src_w / dst_w - 0.5f;
      int x0 = std::max(0, (int)std::floor(fx));
      int x1 = std::min(src_w - 1, x0 + 1);
      float wx = fx - x0;

      for (int c = 0; c < 3; c++) {
        float v00 = src[(y0 * src_w + x0) * 3 + c];
        float v01 = src[(y0 * src_w + x1) * 3 + c];
        float v10 = src[(y1 * src_w + x0) * 3 + c];
        float v11 = src[(y1 * src_w + x1) * 3 + c];
        float v = (1 - wy) * ((1 - wx) * v00 + wx * v01)
                + wy * ((1 - wx) * v10 + wx * v11);
        dst[(y * dst_w + x) * 3 + c] = (uint8_t)v;
      }
    }
  }
}

void frame_preprocess_ffmpeg(const AVFrame* frame,
                             std::size_t img_size,
                             ov::Tensor& input_tensor) {
  int src_w = frame->width;
  int src_h = frame->height;
  const uint8_t* src_data = frame->data[0];  // RGB24
  int src_linesize = frame->linesize[0];

  // 1. Resize (keep aspect ratio, long side = img_size*1.14)
  int target_size = static_cast<int>(img_size * 1.14);
  int new_w, new_h;
  if (src_h < src_w) {
    new_h = target_size;
    new_w = src_w * target_size / src_h;
  } else {
    new_w = target_size;
    new_h = src_h * target_size / src_w;
  }

  std::vector<uint8_t> resized(new_w * new_h * 3);

  // Flatten source to contiguous buffer if needed
  std::vector<uint8_t> src_buf(src_h * src_w * 3);
  for (int y = 0; y < src_h; y++) {
    std::memcpy(src_buf.data() + y * src_w * 3,
                src_data + y * src_linesize,
                src_w * 3);
  }

  bilinear_resize(src_buf.data(), src_w, src_h, resized.data(), new_w, new_h);

  // 2. Center crop img_size x img_size
  int start_x = (new_w - img_size) / 2;
  int start_y = (new_h - img_size) / 2;

  float* tensor_data = input_tensor.data<float>();
  std::vector<float> mean = {0.485f, 0.456f, 0.406f};
  std::vector<float> std  = {0.229f, 0.224f, 0.225f};

  size_t channel_size = img_size * img_size;

  // 3. HWC -> CHW + normalize
  for (size_t c = 0; c < 3; c++) {
    for (size_t y = 0; y < img_size; y++) {
      for (size_t x = 0; x < img_size; x++) {
        int src_x = start_x + x;
        int src_y = start_y + y;
        uint8_t v = resized[(src_y * new_w + src_x) * 3 + c];
        tensor_data[c * channel_size + y * img_size + x] =
          ((float)v / 255.0f - mean[c]) / std[c];
      }
    }
  }
}

cv::Mat avframe_to_cvmat(AVFrame* frame) {
  if (!frame) throw std::runtime_error("AVFrame is null");

  int width = frame->width;
  int height = frame->height;

  cv::Mat img(height, width, CV_8UC3);

  SwsContext* sws_ctx = sws_getContext(
    width, height, (AVPixelFormat)frame->format, // 输入格式
    width, height, AV_PIX_FMT_BGR24,      // 输出 BGR24
    SWS_BILINEAR, nullptr, nullptr, nullptr
  );

  if (!sws_ctx) throw std::runtime_error("Failed to create SwsContext");

  uint8_t* dst_data[4] = { img.data, nullptr, nullptr, nullptr };
  int dst_linesize[4] = { static_cast<int>(img.step), 0, 0, 0 };

  sws_scale(sws_ctx, frame->data, frame->linesize, 0, height, dst_data, dst_linesize);

  sws_freeContext(sws_ctx);
  return img;
}

namespace fs = std::filesystem;

// Save cv::Mat to a file with auto directory creation and conflict resolution
bool save_cvmat(std::string const& path, cv::Mat const& img) {
  try {
    fs::path p(path);

    // ensure parent directory exists
    if (p.has_parent_path()) {
      fs::create_directories(p.parent_path());
    }

    // if file exists, add suffix
    fs::path stem = p.stem();     // filename without extension
    fs::path ext  = p.extension(); // ".jpg", ".png", ...
    fs::path parent = p.parent_path();

    int counter = 1;
    fs::path finalPath = p;
    while (fs::exists(finalPath)) {
      std::ostringstream oss;
      oss << stem.string() << "_" << counter << ext.string();
      finalPath = parent / oss.str();
      counter++;
    }

    return cv::imwrite(finalPath.string(), img);
  } catch (const std::exception& e) {
    std::cerr << "Error saving image: " << e.what() << std::endl;
    return false;
  }
}

thread_local TimeMeasureSpan::MeasureInfo TimeMeasureSpan::time_measure_info;
